import { useEffect, useRef } from "react";
import { getInputs } from "./data";
import * as tf from "@tensorflow/tfjs";
import * as tfvis from "@tensorflow/tfjs-vis";
import * as mobilenet from "@tensorflow-models/mobilenet";
import { img2x, file2img } from "./utils";
const BRAND_CLASSES = [ 'android', 'apple', 'windows' ]

export default function Brand() {
    const truncatedMobilenetRef = useRef(null);
    const modelRef = useRef(null);
    async function start() {
        const { inputs, labels } = await getInputs();
        const surface = tfvis.visor().surface({ name: "输入示例", styles: { height: 250 } })
        inputs.forEach(imgEl => {
            surface.drawArea.appendChild(imgEl);
        })
        const mobilenet = await tf.loadLayersModel("https://storage.googleapis.com/tfjs-models/tfjs/mobilenet_v1_0.25_224/model.json");
        mobilenet.summary();
        const layer = mobilenet.getLayer("conv_pw_13_relu");
        const truncatedMobilenet = tf.model({
            inputs: mobilenet.inputs,
            outputs: layer.output
        });
        const model = tf.sequential();
        model.add(tf.layers.flatten({
            inputShape: layer.outputShape.slice(1),
        }));
        model.add(tf.layers.dense({
            units: 10,
            activation: "relu"
        }))
        model.add(tf.layers.dense({
            units: 3,
            activation: "softmax"
        }))
        model.compile({ loss: "categoricalCrossentropy", optimizer: tf.train.adam() });
        const { xs, ys } = tf.tidy(() => {
            const xs = tf.concat(inputs.map(imgEl => truncatedMobilenet.predict(img2x(imgEl))));
            const ys = tf.tensor(labels);
            return { xs, ys }
        });
        model.fit(xs, ys, {
            epochs: 20,
            callbacks: tfvis.show.fitCallbacks(
                { name: "训练效果" },
                ['loss'],
                { callbacks: ["onEpochEnd"] }
            )
        });
        truncatedMobilenetRef.current = truncatedMobilenet;
        modelRef.current = model;
    }
    async function predict (file) {
        const img = await file2img(file);
        document.body.appendChild(img);
        const pred = tf.tidy(() => {
            const x = img2x(img);
            const input = truncatedMobilenetRef.current.predict(x);
            return modelRef.current.predict(input);
        });
        const index = pred.argMax(1).dataSync()[0];
        setTimeout(() => {
            console.log(`预测结果: ${BRAND_CLASSES[index]}`)
        })
    }
    useEffect(() => {
        start();
    }, [])
    return <div>
        <input type="file" onChange={(e) => {
            predict(e.target.files[0]);
        }} ></input>
    </div>
}